import os

import bs4 as bs4
from dotenv import load_dotenv
from langchain.chains.summarize import load_summarize_chain
from langchain.chains.combine_documents.stuff import StuffDocumentsChain, LLMChain
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import WebBaseLoader

load_dotenv()


# 1.创建模型
model = ChatOpenAI(
    model='qwen-plus',
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 通过DocumentLoader加载数据
loader = WebBaseLoader(
    web_paths=['https://lilianweng.github.io/posts/2023-06-23-agent/'],
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(class_=('post-header', 'post-title', 'post-content'))
    )
)

docs = loader.load()
# print(docs)

# Stuff的第一种写法
# chain = load_summarize_chain(model, chain_type='stuff')
# result = chain.invoke(docs)


# Stuff的第二种写法
prompt_template = """针对以下内容写一个简洁的总结摘要：
"{text}"
简洁的总结摘要：
"""
prompt = PromptTemplate.from_template(prompt_template)

llm_chain = LLMChain(llm=model, prompt=prompt)
stuff_chain = StuffDocumentsChain(llm_chain=llm_chain, document_variable_name="text")
result = stuff_chain.invoke(docs)

print(result['output_text'])
